# helper Module that adds positional encoding to the token embedding to introduce a notion of word order.
import math

import torch
from torch import nn, Tensor
from torch.nn import Transformer
import torch.nn.functional as F


class PositionalEncoding(nn.Module):
    def __init__(self,
                 emb_size: int,
                 dropout: float,
                 max_len: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        pos = torch.arange(0, max_len).reshape(max_len, 1)
        pos_embedding = torch.zeros((max_len, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])


# Seq2Seq Network
def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz))) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask


class Seq2SeqTransformer(nn.Module):
    def __init__(self,
                 num_encoder_layers: int,
                 num_decoder_layers: int,
                 emb_size: int,
                 nhead: int,
                 tgt_vocab_size: int,
                 dim_feedforward: int = 128,
                 dropout: float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        self.transformer = Transformer(d_model=emb_size,
                                       nhead=nhead,
                                       num_encoder_layers=num_encoder_layers,
                                       num_decoder_layers=num_decoder_layers,
                                       dim_feedforward=dim_feedforward,
                                       dropout=dropout,
                                       batch_first=True)
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        self.positional_encoding = PositionalEncoding(
            emb_size, dropout=dropout)

    def forward(self,
                src_emb: Tensor,
                tgt_emb: Tensor,
                src_mask: Tensor = None,
                tgt_mask: Tensor = None,
                src_padding_mask: Tensor = None,
                tgt_padding_mask: Tensor = None,
                memory_key_padding_mask: Tensor = None):
        src_emb = self.positional_encoding(src_emb)
        tgt_emb = self.positional_encoding(tgt_emb)
        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
                                src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
        outs = self.generator(outs)
        outs = F.log_softmax(outs, dim=-1)
        return outs

    def encode(self, src_emb: Tensor, src_mask: Tensor = None):
        return self.transformer.encoder(self.positional_encoding(
            src_emb), src_mask)

    def decode(self, tgt_emb: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer.decoder(self.positional_encoding(
            tgt_emb), memory, tgt_mask)

    # function to generate output sequence using greedy algorithm
    def greedy_decode(self, src, tgt_embedded, start_symbol, end_symbol, max_len=20, device="cpu"):
        src = src.to(device)
        memory = self.encode(src)

        ys = torch.ones(src.shape, 1).fill_(start_symbol).type(torch.long).to(device)
        ys_embedded = tgt_embedded[ys]
        ys_embedded.transpose(0, 1)

        decoder_outputs = []

        for i in range(max_len - 1):
            memory = memory.to(device)
            tgt_mask = (generate_square_subsequent_mask(ys.size(0))
                        .type(torch.bool)).to(device)
            out = self.decode(ys_embedded, memory, tgt_mask)
            prob = self.generator(out[:, -1])
            decoder_output = self.generator(out[:, -1:])
            decoder_outputs.append(decoder_output)

            _, next_word = torch.max(prob, dim=1)
            next_word = next_word.item()

            ys = torch.cat([ys, torch.ones(1, 1).fill_(next_word).type(torch.long)], dim=0)
            ys_embedded = tgt_embedded[ys]
            ys_embedded = ys_embedded.transpose(0, 1)

            if next_word == end_symbol:
                break

        decoder_outputs = torch.cat(decoder_outputs, dim=1)
        decoder_outputs = F.log_softmax(decoder_outputs, dim=-1)

        return decoder_outputs, ys
